#!/usr/bin/env python3

from functools import partial

import lhotse
from lhotse import Fbank, FbankConfig
from lhotse.dataset import SimpleCutSampler
from lhotse.dataset.input_strategies import (
    AudioSamples,
    OnTheFlyFeatures,
)
from torch.utils.data.dataloader import DataLoader

from speech_recognition_dataset import ConsistencyRegularizationSpeechRecognitionDataset


def create_cutset():
    recording0 = lhotse.Recording.from_file("./0.wav")
    sup0 = lhotse.SupervisionSegment(
        id="sup0",
        recording_id=recording0.id,
        start=0,
        duration=5,
        text="hello, how are you",
    )
    cut0 = lhotse.MonoCut(
        id="cut0",
        start=0,
        duration=6,
        channel=0,
        recording=recording0,
        supervisions=[sup0],
    )

    recording1 = lhotse.Recording.from_file("./1.wav")
    sup1 = lhotse.SupervisionSegment(
        id="sup1",
        recording_id=recording1.id,
        start=0,
        duration=3,
        text="fine, thank you",
    )
    cut1 = lhotse.MonoCut(
        id="cut1",
        start=0,
        duration=4,
        channel=0,
        recording=recording1,
        supervisions=[sup1],
    )

    return lhotse.CutSet([cut0, cut1])


def main():
    cutset = create_cutset()
    print([c for c in cutset])

    t1 = partial(lhotse.MonoCut.perturb_speed, factor=0.9)
    t2 = partial(lhotse.MonoCut.perturb_volume, factor=1.1)
    t3 = partial(lhotse.MonoCut.perturb_tempo, factor=1.2)
    transforms = [t1, t2, t3]

    train = ConsistencyRegularizationSpeechRecognitionDataset(
        input_strategy=OnTheFlyFeatures(Fbank(FbankConfig(num_mel_bins=80))),
        #  input_strategy=AudioSamples(),
        cut_transforms=transforms,
        #  return_cuts=True,
        return_cuts=False,
    )
    print(train)
    train_sampler = SimpleCutSampler(
        cutset,
        max_duration=20,
        shuffle=False,
    )

    train_dl = DataLoader(
        train,
        sampler=train_sampler,
        batch_size=None,
        num_workers=1,
        persistent_workers=False,
    )
    for b in train_dl:
        print(b["inputs"].shape, b["supervisions"])
        if "aug" in b:
            assert len(b["aug"]) == len(transforms)
            for i, aug in enumerate(b["aug"]):
                print(
                    transforms[i].func.__name__,
                    aug["inputs"].shape,
                    aug["supervisions"],
                )


if __name__ == "__main__":
    main()